# Author: Wang Chenxu
# Last Edit: Wang Chenxu
# Create Data: 2023.12.01 19:31 PM
# Last Edit Data: 2023.12.01 19:31 PM
# Thanks: Xu YuCheng

import nuke
import math
# Initiate node everytime you hit split button
# Removing all node in nodeClass list
def clearNodes():
    nodeClass = ['Input', 'Expression', 'DeepRecolor', 'Remove']
    for n in nuke.allNodes():
        if n.Class() in nodeClass:
            nuke.delete(n)

# Deselect everything
def deselect():
    for n in nuke.allNodes():
        n.setSelected(False)

# Get the layer we don't want, and put them into Remove node after the dot node
def getOtherLayer(chList, nameList):
    layerList = [i.split('.')[0] for i in chList]
    finalList = list(set(layerList))
    nameList.append('rgba')
    for element in nameList:
        finalList.remove(element)
    return finalList

# Main program
def deepLayer():
    tn = nuke.thisNode()
    tn.end()

    snList = nuke.selectedNodes()

    # if none is selected or selecting itself, it will only clear up the scene and ready to dance
    if len(snList) == 0 or (snList[0] is tn):    	
        inputAccount = tn.inputs()
        for input in range(inputAccount):
            tn.setInput(input, None)
        tn.begin()
        clearNodes()
        tn.end()

    else:
        tn.begin()
        clearNodes()

        # Prepare node for calling
        blackPlate = nuke.toNode('blackPlate')
        deepMerge = nuke.toNode('DeepMerge')
        finalDot = nuke.toNode('Dot1')
        output = nuke.toNode('Output1')

        # inGroupCounter is for calculating the input of deepMerge
        inGroupCounter = 0
        for n in snList:

            # Create layers we need
            name = n['label'].getValue()
            nuke.Layer(name, ['{0}.alpha'.format(name)])

        	# Create input for each deep input data
            inputNode = nuke.createNode('Input', inpanel=False)
            inputNode['name'].setValue(name)
            deselect()

            # Create DeepRecolor node
            deepRecolorNode = nuke.createNode('DeepRecolor', inpanel=False)
            deepRecolorNode['name'].setValue(name + '_Recolor')
            deepRecolorNode['channels'].setValue(name)
            deepRecolorNode.setInput(0, inputNode)
            deselect()

            # Create expression node to define a alpha channel for specific layer
            expressionNode = nuke.createNode('Expression', inpanel=False)
            expressionNode['name'].setValue(name + '_Expression')
            expressionNode['channel0'].setValue(name)
            expressionNode['expr0'].setValue('1')
            expressionNode.setInput(0, blackPlate)
            deepRecolorNode.setInput(1, expressionNode)
            deselect()

            deepMerge.setInput(inGroupCounter, deepRecolorNode)
            inGroupCounter += 1

        tn.end()

        counter = 0
        for n in snList:
            tn.setInput(counter, n)
            counter += 1

# Clear other layers but specific layers we need
# The ideas is collecting layers of inputs node have, and collect the name of specific layers 
# Compare those two list and delete those layers we need from the inputs layers
# Then use Remove node to remove them
def clearChannel():
    tn = nuke.thisNode()
    tn.end()

    inputAccount = tn.inputs()
    snList = []
    for i in range(inputAccount):
        snList.append(tn.input(i))

    tn.begin()

    for n in nuke.allNodes():
        if n.Class() in ['Remove']:
            nuke.delete(n)

    finalDot = nuke.toNode('Dot1')
    output = nuke.toNode('Output1')

    nameList = [n['label'].getValue() for n in snList]
    ch = finalDot.channels()
    channelList = getOtherLayer(ch, nameList)

    channelLen = len(channelList)
    lastRemoveNode = ''

    for i in range(math.ceil(float(channelLen) / 4)):
        channelAccount = 1
        removeNode = nuke.createNode('Remove', inpanel=False)

        deselect()

        for j in range(4):
            if channelAccount == 1:
                indexs = ''
            else:
                indexs = str(channelAccount)
            channelAccount += 1
            try:
                removeNode['channels{0}'.format(indexs)].setValue(channelList[i * 4 + j])
            except Exception as e:
                pass

        if i == 0:
            removeNode.setInput(0, finalDot)
        else:
            removeNode.setInput(0, lastRemoveNode)

        lastRemoveNode = removeNode
        
    output.setInput(0, removeNode)